from PIL import Image

from hhutil.io import fmt_path, eglob

from paddle.fluid.io import Dataset


def make_divisable(size, factor=2):
    w, h = size
    w = w // factor * factor
    h = h // factor * factor
    return w, h


class SRTest(Dataset):

    def __init__(self, root, scale, ext='png', transform=None):
        super().__init__()
        self.root = fmt_path(root)
        self.scale = scale
        self.ext = ext
        self.transform = transform

        self.filenames = list(eglob(self.root, "*." + self.ext))

    def __getitem__(self, item):
        hr = Image.open(self.filenames[item]).convert('RGB')
        w, h = make_divisable(hr.size, self.scale)
        hr = hr.crop([0, 0, w, h])
        lr = hr.resize((w // self.scale, h // self.scale), Image.BICUBIC)
        if self.transform is not None:
            lr, hr = self.transform(lr, hr)
        return lr, hr

    def __len__(self):
        return len(self.filenames)